import os
import platform
import shutil
from datetime import datetime as dt
from common import get_images, remove_imgs


ROOT_DIR = os.path.abspath('./')
ASCENDBASE_DIR = os.path.abspath("./ascend-docker-image/Dockerfile/latest/common/")
ALGORITHM_DIR = os.path.abspath("./ascend-docker-image/Dockerfile/latest/common/ascend-algorithm")


class Logger():
    def __init__(self, file_name) -> None:
        self.file_name = file_name
    
    def info(self, msg):
        t = dt.now().strftime("%Y-%m-%d %H:%M:%S")
        with open(self.file_name, 'a') as f:
            f.write(t + "[INFO] " + str(msg) + '\n')
    
    def error(self, msg):
        t = dt.now().strftime("%Y-%m-%d %H:%M:%S")
        with open(self.file_name, 'a') as f:
            f.write(t + "[ERROR] " + str(msg) + '\n')

logger = Logger(os.path.join(ROOT_DIR, "docker_build.log"))


class Image():
    def __init__(self, name, tag, _id=None) -> None:
        self.name = name
        self.tag = tag
        self.id = _id
        self.res = ""
    
    def remove(self):
        res = os.popen("docker rmi %s:%s" % (self.name, self.tag)).read()
        if self._check_error(res):
            return False
        return True
    
    @staticmethod
    def _check_error(res, includes=['error', 'failed']):
        for i in includes:
            if i in res.lower():
                return True
        return False


def check_results(res, includes=['error', 'failed']):
    for i in includes:
        if i in res.lower():
            return True
    return False


def prepare_check(need_files):
    listdirs = os.listdir(ROOT_DIR)
    for f in need_files:
        if f not in listdirs:
            raise OSError("%s not exist" % f)


def remove_last_images():
    if platform.system() == "Linux":
        imgs = get_images('docker images |grep algorithm')
    elif platform.system() == "Windows":
        imgs = get_images('docker images |findstr algorithm')
    remove_imgs(imgs)


def build_ascendbase(mode, systems):
    base_dir = os.path.join(ASCENDBASE_DIR, 'ascendbase-' + mode)
    arch_dir = {'aarch64': 'arm64', 'x86_64': 'x64', "AMD64": 'x64'}
    arch = platform.machine()
    docker_arch = {'aarch64': 'arm64', 'x86_64': 'amd64', "AMD64": 'amd64'}
    libstdc_path = os.path.join(ROOT_DIR, 'libstdc++.so.6.0.24')
    for sys_name in systems.keys():
        run_dir = os.path.join(base_dir, sys_name + '-' + arch_dir[arch])
        libstdc_run_file = os.path.join(run_dir, 'libstdc++.so.6.0.24')
        if not os.path.exists(libstdc_run_file):
            shutil.copyfile(libstdc_path, libstdc_run_file)
        os.chdir(run_dir)
        tag = sys_name + '-' + docker_arch[arch]
        cmd = 'docker build -t ascendbase-%s:%s --no-cache=true .' %(mode, tag)
        print('exec %s' % cmd)
        res = os.system(cmd)
        if res != 0:
            print('build ascendbase %s failed' % sys_name)
            logger.error(res)


def remove(file_or_dir):
    if not os.path.exists(file_or_dir):
        return
    if os.path.isfile(file_or_dir):
        os.remove(file_or_dir)
        return
    elif os.path.isdir(file_or_dir):
        shutil.rmtree(file_or_dir, ignore_errors=True)


def build_algorithm(systems):
    arch = platform.machine()
    pkg_arch = {'AMD64': 'x86_64'}
    arch = pkg_arch.get(arch, arch)
    docker_arch = {'aarch64': 'arm64', 'x86_64': 'amd64'}
    os.chdir(ROOT_DIR)
    install_info_file = os.path.join(ALGORITHM_DIR, 'ascend_install.info')
    version_info = os.path.join(ALGORITHM_DIR, 'version.info')
    if not os.path.exists(install_info_file) or not os.path.exists(version_info):
        remove(install_info_file)
        remove(version_info)
        shutil.copyfile('ascend_install.info', install_info_file)
        shutil.copyfile('version.info', version_info)
    for sys_name in systems.keys():
        for cann_ver in systems[sys_name]:
            os.chdir(ROOT_DIR)
            nnrt_pkg = "Ascend-cann-nnrt_%s_linux-%s.run" % (cann_ver, arch)
            nnrt_pkg_path = os.path.join(ALGORITHM_DIR, nnrt_pkg)
            if not os.path.exists(nnrt_pkg_path):
                shutil.copyfile(nnrt_pkg, nnrt_pkg_path)
            base_tag = sys_name + '-' + docker_arch[arch]
            tag = cann_ver + '-' + sys_name + '-' + docker_arch[arch]
            os.chdir(ALGORITHM_DIR)
            cmd = "docker build -t algorithm:%s --build-arg NNRT_PKG=%s --build-arg BASE_VERSION=%s ." % (tag, nnrt_pkg, base_tag)
            res = os.system(cmd)
            if res != 0:
                print('build algrithm %s failed' % tag)
                logger.error('build algrithm %s failed' % tag)
            remove(nnrt_pkg_path)


def upload_algorithm(systems):
    arch = platform.machine()
    docker_arch = {'aarch64': 'arm64', 'x86_64': 'amd64', 'AMD64': 'amd64'}
    for sys_name in systems.keys():
        for cann_ver in systems[sys_name]:
            img = "algorithm:%s-%s-%s" % (cann_ver, sys_name, docker_arch[arch])
            os.system("docker tag %s swr.cn-north-1.myhuaweicloud.com/liuyazhou/%s" % (img, img))
            os.system("docker push swr.cn-north-1.myhuaweicloud.com/liuyazhou/%s" % img)


def update_code():
    os.system(
        "cd ascend-docker-image &&"
        "git clean -f -d&&"
        "git reset --hard HEAD^ &&"
        "git pull"
    )


if __name__ == "__main__":
    need_build_base = True
    base_mode = 'train'
    systems = {
        'centos7.6': ['5.0.1', '20.3.0'],
        'openeuler20.03': ['5.0.2']
    }
    img_name = 'algorithm'
    # run 包放前面，顺序不要变
    arch = platform.machine()
    pkg_arch = {'AMD64': 'x86_64'}
    arch = pkg_arch.get(arch, arch)
    need_files = [
        'Ascend-cann-nnrt_5.0.1_linux-%s.run' % arch,
        'Ascend-cann-nnrt_5.0.2_linux-%s.run' % arch,
        'Ascend-cann-nnrt_20.3.0_linux-%s.run' % arch,
        'ascend_install.info',
        'version.info',
        'libstdc++.so.6.0.24'
    ]
    prepare_check(need_files)
    update_code()
    if need_build_base:
        remove_last_images()
        print('build ascendbase begin....')
        build_ascendbase(base_mode, systems)
    build_algorithm(systems)
    upload_algorithm(systems)
